package com.wzz.utils.kms;

import java.util.*;

public class DecisionTree {
    private Node root;

    // 决策树节点
    private static class Node {
        private final String attribute; // 节点对应的属性
        private final Map<String, Node> children; // 子节点
        private final String result; // 叶节点对应的结果

        // 内部类的构造函数
        private Node(String attribute) {
            this.attribute = attribute;
            this.children = new HashMap<>();
            this.result = null;
        }

        private Node(String result, boolean isLeaf) {
            this.attribute = null;
            this.children = null;
            this.result = result;
        }

        private boolean isLeaf() {
            return result != null;
        }
    }

    // 训练决策树
    public void train(List<Map<String, String>> data) {
        List<String> attributes = new ArrayList<>(data.get(0).keySet()); // 获取数据中所有属性的列表
        attributes.remove("label"); // 移除类别属性

        root = buildTree(data, attributes);
    }

    // 递归地构建决策树
    private Node buildTree(List<Map<String, String>> data, List<String> attributes) {
        if (data.isEmpty()) {
            return new Node(null, true);
        }

        // 所有样本属于同一类别
        String firstClass = data.get(0).get("label");
        boolean allSameClass = data.stream().allMatch(d -> d.get("label").equals(firstClass));
        if (allSameClass) {
            return new Node(firstClass, true);
        }

        // 没有属性可分裂
        if (attributes.isEmpty()) {
            return new Node(majorityClass(data), true);
        }

        // 选择最佳分裂属性
        String splitAttribute = selectBestAttribute(data, attributes);

        Node node = new Node(splitAttribute);

        // 递归构建子节点
        Map<String, List<Map<String, String>>> splits = splitData(data, splitAttribute);
        for (Map.Entry<String, List<Map<String, String>>> entry : splits.entrySet()) {
            String value = entry.getKey();
            List<Map<String, String>> subset = entry.getValue();
            List<String> remainingAttributes = new ArrayList<>(attributes);
            remainingAttributes.remove(splitAttribute);

            Node child = buildTree(subset, remainingAttributes);
            node.children.put(value, child);
        }

        return node;
    }

    // 预测单个样本
    public String predict(Map<String, String> instance) {
        Node node = root;
        while (!node.isLeaf()) {
            String attribute = node.attribute;
            String value = instance.get(attribute);
            node = node.children.get(value);
        }
        return node.result;
    }

    // 获取数据集中类别最多的类
    private String majorityClass(List<Map<String, String>> data) {
        Map<String, Integer> counts = new HashMap<>();
        for (Map<String, String> instance : data) {
            String cls = instance.get("label");
            counts.put(cls, counts.getOrDefault(cls, 0) + 1);
        }
        return Collections.max(counts.entrySet(), Map.Entry.comparingByValue()).getKey();
    }

//
    // 计算信息熵
    private double entropy(List<Map<String, String>> data) {
    int n = data.size();
    if (n <= 1) {
        return 0;
    }

    Map<String, Integer> counts = new HashMap<>();
    for (Map<String, String> instance : data) {
        String cls = instance.get("label");
        counts.put(cls, counts.getOrDefault(cls, 0) + 1);
    }

    double entropy = 0;
    for (int count : counts.values()) {
        double p = (double) count / n;
        entropy -= p * Math.log(p) / Math.log(2);
    }

    return entropy;
}

    // 计算信息增益
    private double informationGain(List<Map<String, String>> data, String attribute) {
        double entropyBeforeSplit = entropy(data);

        Map<String, List<Map<String, String>>> splits = splitData(data, attribute);
        double entropyAfterSplit = 0;
        int n = data.size();
        for (List<Map<String, String>> subset : splits.values()) {
            double p = (double) subset.size() / n;
            entropyAfterSplit += p * entropy(subset);
        }

        return entropyBeforeSplit - entropyAfterSplit;
    }

    // 选择信息增益最大的属性
    private String selectBestAttribute(List<Map<String, String>> data, List<String> attributes) {
        String bestAttribute = null;
        double maxInformationGain = -1;

        for (String attribute : attributes) {
            double informationGain = informationGain(data, attribute);
            if (informationGain > maxInformationGain) {
                maxInformationGain = informationGain;
                bestAttribute = attribute;
            }
        }

        return bestAttribute;
    }

    // 根据属性分裂数据
    private Map<String, List<Map<String, String>>> splitData(List<Map<String, String>> data, String attribute) {
        Map<String, List<Map<String, String>>> splits = new HashMap<>();
        for (Map<String, String> instance : data) {
            String value = instance.get(attribute);
            List<Map<String, String>> subset = splits.getOrDefault(value, new ArrayList<>());
            subset.add(instance);
            splits.put(value, subset);
        }

        return splits;
    }

  public static void main(String[] args) {
      // 训练集
      List<Map<String, String>> data = new ArrayList<>();
      // 添加训练数据
      Map<String, String> instance1 = new HashMap<>();
      instance1.put("outlook", "sunny");
      instance1.put("temperature", "hot");
      instance1.put("humidity", "high");
      instance1.put("windy", "false");
      instance1.put("label", "no");
      instance1.put("label", "no2");
      instance1.put("label", "no3");
      data.add(instance1);
     //自己编数据吧。。。
     //构建决策树
      DecisionTree decisionTree = new DecisionTree();
      decisionTree.train(data);
      // 进行预测
      Map<String, String> instance = new HashMap<>();
      instance.put("outlook", "sunny");
      instance.put("temperature", "cool");
      instance.put("humidity", "high");
      instance.put("windy", "false");
      String prediction = decisionTree.predict(instance);
      System.out.println("Prediction: " + prediction);
  }
}
